#!/bin/bash

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
NC='\033[0m' # 重置颜色

# 配置信息
CONFIG_FILE="$HOME/.ssh_proxy_config"
LOG_FILE="$HOME/ssh_proxy.log"
PROXYCHAINS_CONF="/etc/proxychains4.conf"
PROXYCHAINS_CONF_BAK="/etc/proxychains4.conf.bak"
REDSOCKS_CONF="/etc/redsocks.conf"
REDSOCKS_CONF_BAK="/etc/redsocks.conf.bak"

# 日志函数：记录操作日志并输出到终端
log() {
    local level="$1"
    local message="$2"
    local timestamp=$(date +"%Y-%m-%d %H:%M:%S")
    echo "[$timestamp] [$level] $message" | tee -a "$LOG_FILE"
}

# 检查命令是否存在（优化：增加版本检查基础逻辑）
check_command() {
    command -v "$1" >/dev/null 2>&1
}

# 检查是否有root权限
check_root() {
    if [ "$(id -u)" -ne 0 ]; then
        log "ERROR" "请使用root权限运行此脚本"
        echo -e "${RED}错误：此脚本需要root权限，请使用sudo运行${NC}"
        exit 1
    fi
}

# 检查依赖是否安装（按系统分类处理）
check_dependencies() {
    local missing_deps=""
    
    # 根据包管理器定义依赖列表
    if check_command apt-get; then
        # Debian/Ubuntu系列
        DEPENDENCIES="proxychains4 redsocks iptables lsof curl"
    elif check_command yum; then
        # CentOS/RHEL系列
        DEPENDENCIES="openssh-clients proxychains-ng redsocks iptables lsof curl"
    elif check_command pacman; then
        # Arch Linux系列
        DEPENDENCIES="openssh proxychains-ng redsocks iptables lsof curl"
    else
        log "ERROR" "无法识别包管理器"
        echo -e "${RED}错误：无法识别系统包管理器，请手动安装依赖${NC}"
        exit 1
    fi
    
    # 检查依赖是否缺失
    for dep in $DEPENDENCIES; do
        if ! check_command "$dep"; then
            missing_deps="$missing_deps $dep"
        fi
    done
    
    echo "$missing_deps"
}

# 安装缺失的依赖（优化：增加安装后二次检查）
install_dependencies() {
    local missing_deps=$(check_dependencies)
    
    if [ -z "$missing_deps" ]; then
        log "INFO" "所有依赖已安装"
        return 0
    fi
    
    log "INFO" "发现缺失的依赖: $missing_deps，开始安装..."
    echo -e "${BLUE}正在安装缺失的依赖: $missing_deps${NC}"
    
    # 检测包管理器并执行安装
    if check_command apt-get; then
        # Debian/Ubuntu
        sudo apt-get update >/dev/null 2>&1
        sudo apt-get install -y $missing_deps >/dev/null 2>&1
    elif check_command yum; then
        # CentOS/RHEL
        sudo yum install -y $missing_deps >/dev/null 2>&1
    elif check_command pacman; then
        # Arch Linux
        sudo pacman -Sy --noconfirm $missing_deps >/dev/null 2>&1
    else
        log "ERROR" "无法识别包管理器，安装依赖失败"
        echo -e "${RED}错误：无法安装依赖，请手动安装: $missing_deps${NC}"
        exit 1
    fi
    
    # 二次检查是否安装成功（针对关键依赖redsocks）
    local failed_deps=""
    for dep in $missing_deps; do
        if ! check_command "$dep"; then
            failed_deps="$failed_deps $dep"
        fi
    done
    
    if [ -n "$failed_deps" ]; then
        log "ERROR" "依赖安装失败: $failed_deps"
        echo -e "${RED}错误：依赖安装失败，请手动安装: $failed_deps${NC}"
        # 特别提示redsocks的备用安装方案
        if echo "$failed_deps" | grep -q "redsocks"; then
            echo -e "${YELLOW}提示：redsocks可尝试源码安装：https://github.com/darkk/redsocks${NC}"
        fi
        exit 1
    else
        log "INFO" "依赖安装成功"
        echo -e "${GREEN}依赖安装成功！${NC}"
    fi
}

# 查找指定范围内的可用端口
find_available_port() {
    local start_port="$1"
    local end_port="$2"
    local port="$start_port"
    
    while [ "$port" -le "$end_port" ]; do
        if ! lsof -Pi :$port -sTCP:LISTEN -t >/dev/null; then
            echo "$port"
            return 0
        fi
        ((port++))
    done
    
    log "ERROR" "在$start_port-$end_port范围内找不到可用端口"
    echo -e "${RED}错误：端口范围$start_port-$end_port内无可用端口${NC}"
    exit 1
}

# 配置proxychains代理
configure_proxychains() {
    local proxy_port="$1"
    log "INFO" "配置proxychains使用端口$proxy_port"
    
    # 备份原始配置（仅首次操作）
    if [ ! -f "$PROXYCHAINS_CONF_BAK" ]; then
        sudo cp "$PROXYCHAINS_CONF" "$PROXYCHAINS_CONF_BAK"
        log "INFO" "已备份原始proxychains配置"
    fi
    
    # 写入新配置
    sudo cat > "$PROXYCHAINS_CONF" << EOF
strict_chain
proxy_dns 
remote_dns_subnet 224
tcp_read_time_out 15000
tcp_connect_time_out 8000
[ProxyList]
socks5  127.0.0.1 $proxy_port
EOF
    
    log "INFO" "proxychains配置完成"
}

# 配置redsocks透明代理
configure_redsocks() {
    local proxy_port="$1"
    local redsocks_port="$2"
    
    log "INFO" "配置redsocks，监听端口$redsocks_port，转发到$proxy_port"
    
    # 备份原始配置（仅首次操作）
    if [ ! -f "$REDSOCKS_CONF_BAK" ]; then
        sudo cp "$REDSOCKS_CONF" "$REDSOCKS_CONF_BAK" 2>/dev/null || true
        log "INFO" "已备份原始redsocks配置"
    fi
    
    # 写入新配置
    sudo cat > "$REDSOCKS_CONF" << EOF
base {
    log_info = on;
    log = "file:/var/log/redsocks.log";
    daemon = on;
    user = redsocks;
    group = redsocks;
    redirector = iptables;
}

redsocks {
    local_ip = 127.0.0.1;
    local_port = $redsocks_port;
    ip = 127.0.0.1;
    port = $proxy_port;
    type = socks5;
}
EOF
    
    log "INFO" "redsocks配置完成"
}

# 配置系统代理环境变量
configure_system_proxy() {
    local proxy_port="$1"
    
    log "INFO" "配置系统代理环境变量"
    
    # 设置当前会话代理
    export http_proxy="socks5h://127.0.0.1:$proxy_port"
    export https_proxy="socks5h://127.0.0.1:$proxy_port"
    export ftp_proxy="socks5h://127.0.0.1:$proxy_port"
    export all_proxy="socks5h://127.0.0.1:$proxy_port"
    export no_proxy="localhost,127.0.0.1,localaddress,.localdomain.com"
    
    # 保存代理设置到用户文件
    cat > "$HOME/.proxyrc" << EOF
export http_proxy="socks5h://127.0.0.1:$proxy_port"
export https_proxy="socks5h://127.0.0.1:$proxy_port"
export ftp_proxy="socks5h://127.0.0.1:$proxy_port"
export all_proxy="socks5h://127.0.0.1:$proxy_port"
export no_proxy="localhost,127.0.0.1,localaddress,.localdomain.com"
EOF
    
    log "INFO" "系统代理环境变量配置完成"
}

# 清除系统代理环境变量
clear_system_proxy() {
    log "INFO" "清除系统代理环境变量"
    
    # 清除当前会话代理
    unset http_proxy
    unset https_proxy
    unset ftp_proxy
    unset all_proxy
    unset no_proxy
    
    # 删除代理配置文件
    if [ -f "$HOME/.proxyrc" ]; then
        rm -f "$HOME/.proxyrc"
    fi
    
    log "INFO" "系统代理环境变量已清除"
}

# 配置iptables规则（用于全局代理，优化：增加持久化提示）
configure_iptables() {
    local redsocks_port="$1"
    
    log "INFO" "配置iptables规则，重定向流量到redsocks端口$redsocks_port"
    
    # 创建自定义链
    iptables -t nat -N REDSOCKS 2>/dev/null
    iptables -t nat -N REDSOCKS2 2>/dev/null
    
    # 清空链规则（防止重复添加）
    iptables -t nat -F REDSOCKS
    iptables -t nat -F REDSOCKS2
    
    # 排除本地和内网流量
    iptables -t nat -A REDSOCKS -d 0.0.0.0/8 -j RETURN
    iptables -t nat -A REDSOCKS -d 10.0.0.0/8 -j RETURN
    iptables -t nat -A REDSOCKS -d 127.0.0.0/8 -j RETURN
    iptables -t nat -A REDSOCKS -d 169.254.0.0/16 -j RETURN
    iptables -t nat -A REDSOCKS -d 172.16.0.0/12 -j RETURN
    iptables -t nat -A REDSOCKS -d 192.168.0.0/16 -j RETURN
    iptables -t nat -A REDSOCKS -d 224.0.0.0/4 -j RETURN
    iptables -t nat -A REDSOCKS -d 240.0.0.0/4 -j RETURN
    
    # 排除SSH服务器自身流量（避免死循环）
    if [ -n "$SSH_SERVER_IP" ]; then
        iptables -t nat -A REDSOCKS -d "$SSH_SERVER_IP" -j RETURN
    fi
    
    # 重定向其他TCP流量到redsocks
    iptables -t nat -A REDSOCKS -p tcp -j REDIRECT --to-ports "$redsocks_port"
    
    # 应用到本地输出流量
    iptables -t nat -A REDSOCKS2 -p tcp -m owner --uid-owner proxy -j RETURN
    iptables -t nat -A REDSOCKS2 -p tcp -m owner --gid-owner proxy -j RETURN
    iptables -t nat -A REDSOCKS2 -j REDSOCKS
    iptables -t nat -A OUTPUT -p tcp -j REDSOCKS2 2>/dev/null
    
    # 应用到转发流量
    iptables -t nat -A PREROUTING -p tcp -j REDSOCKS 2>/dev/null
    
    log "INFO" "iptables规则配置完成"
    echo -e "${YELLOW}提示：若需重启后保留iptables规则，可执行 'iptables-save > /etc/iptables.rules' 并配置开机加载${NC}"
}

# 清除iptables规则
clear_iptables() {
    log "INFO" "清除iptables规则"
    
    # 清除自定义链
    iptables -t nat -F REDSOCKS 2>/dev/null
    iptables -t nat -F REDSOCKS2 2>/dev/null
    iptables -t nat -X REDSOCKS 2>/dev/null
    iptables -t nat -X REDSOCKS2 2>/dev/null
    
    # 移除链引用
    iptables -t nat -D OUTPUT -p tcp -j REDSOCKS2 2>/dev/null
    iptables -t nat -D PREROUTING -p tcp -j REDSOCKS 2>/dev/null
    
    log "INFO" "iptables规则已清除"
}

# 保存配置到文件
save_config() {
    local host="$1"
    local port="$2"
    local user="$3"
    local proxy_port="$4"
    local redsocks_port="$5"
    local ssh_server_ip="$6"
    local proxy_mode="$7"
    
    cat > "$CONFIG_FILE" << EOF
HOST=$host
PORT=$port
USER=$user
PROXY_PORT=$proxy_port
REDSOCKS_PORT=$redsocks_port
SSH_SERVER_IP=$ssh_server_ip
PROXY_MODE=$proxy_mode
EOF
    
    log "INFO" "配置已保存到$CONFIG_FILE"
}

# 加载配置文件
load_config() {
    if [ -f "$CONFIG_FILE" ]; then
        source "$CONFIG_FILE"
        log "INFO" "已加载配置: $USER@$HOST:$PORT，代理端口: $PROXY_PORT，模式: $PROXY_MODE"
        return 0
    else
        log "INFO" "未找到配置文件$CONFIG_FILE"
        return 1
    fi
}

# 启动SSH隧道（优化：保留错误日志，精确检查运行状态）
start_tunnel() {
    local host="$1"
    local port="$2"
    local user="$3"
    local proxy_port="$4"
    
    log "INFO" "启动SSH隧道: $user@$host:$port，本地代理端口$proxy_port"
    
    # 精确检查隧道是否已运行（通过端口监听判断）
    if ss -ltn | grep -q ":$proxy_port "; then
        log "INFO" "SSH隧道已在运行（端口$proxy_port已监听）"
        return 0
    fi
    
    # 启动SSH动态端口转发（保留错误输出到日志）
    ssh -fNTD "$proxy_port" "$user@$host" -p "$port" 2>> "$LOG_FILE"
    
    # 检查启动结果
    if [ $? -ne 0 ]; then
        log "ERROR" "SSH隧道启动失败（查看日志$LOG_FILE获取详细错误）"
        echo -e "${RED}错误：SSH隧道连接失败，可能原因：服务器不可达/认证失败/端口被占用${NC}"
        echo -e "${YELLOW}提示：手动测试连接命令：ssh -vNT $user@$host -p $port${NC}"
        return 1
    fi
    
    # 等待隧道初始化
    sleep 2
    if ss -ltn | grep -q ":$proxy_port "; then
        log "INFO" "SSH隧道启动成功（端口$proxy_port已监听）"
        echo -e "${GREEN}SSH隧道启动成功！${NC}"
        return 0
    else
        log "ERROR" "SSH隧道启动后未正常监听端口$proxy_port"
        echo -e "${RED}错误：SSH隧道启动失败（端口未监听）${NC}"
        return 1
    fi
}

# 停止SSH隧道
stop_tunnel() {
    local proxy_port="$1"
    
    log "INFO" "停止SSH隧道，代理端口$proxy_port"
    
    # 查找并终止隧道进程（通过端口关联的SSH进程）
    local pids=$(lsof -i :$proxy_port -sTCP:LISTEN -t)
    if [ -n "$pids" ]; then
        kill -TERM $pids >/dev/null 2>&1
        # 等待进程终止
        sleep 2
        if lsof -i :$proxy_port -sTCP:LISTEN -t >/dev/null; then
            log "WARNING" "SSH隧道进程无法正常终止，强制杀死"
            kill -KILL $pids >/dev/null 2>&1
        fi
        log "INFO" "SSH隧道已停止"
    else
        log "INFO" "未找到运行中的SSH隧道（端口$proxy_port未监听）"
    fi
}

# 启动redsocks服务（用于全局代理）
start_redsocks() {
    local redsocks_port="$1"
    
    log "INFO" "启动redsocks服务，监听端口$redsocks_port"
    
    # 检查服务是否已运行
    if pgrep -f "redsocks -c $REDSOCKS_CONF" >/dev/null; then
        log "INFO" "redsocks已在运行"
        return 0
    fi
    
    # 启动redsocks
    redsocks -c "$REDSOCKS_CONF" >/dev/null 2>&1
    
    # 检查启动结果
    if [ $? -ne 0 ] || ! pgrep -f "redsocks -c $REDSOCKS_CONF" >/dev/null; then
        log "ERROR" "redsocks启动失败"
        echo -e "${RED}错误：redsocks启动失败${NC}"
        return 1
    fi
    
    log "INFO" "redsocks启动成功"
}

# 停止redsocks服务
stop_redsocks() {
    log "INFO" "停止redsocks服务"
    
    # 查找并终止进程
    local pids=$(pgrep -f "redsocks -c $REDSOCKS_CONF")
    if [ -n "$pids" ]; then
        kill $pids >/dev/null 2>&1
        log "INFO" "redsocks已停止"
    else
        log "INFO" "未找到运行中的redsocks进程"
    fi
}

# 测试代理连接（优化：使用多个测试地址）
test_proxy() {
    local proxy_port="$1"
    local proxy_mode="$2"
    
    log "INFO" "测试代理连接，模式: $proxy_mode"
    
    # 多个测试地址（提高可靠性）
    local test_urls=("https://ipinfo.io/ip" "https://icanhazip.com" "https://ifconfig.me")
    local result=""
    
    if [ "$proxy_mode" = "global" ]; then
        echo -e "${BLUE}测试全局代理连接...${NC}"
        for url in "${test_urls[@]}"; do
            result=$(curl -s --connect-timeout 5 "$url")
            if [ -n "$result" ]; then
                break
            fi
        done
    else
        echo -e "${BLUE}测试proxychains代理连接...${NC}"
        for url in "${test_urls[@]}"; do
            result=$(proxychains4 curl -s --connect-timeout 5 "$url")
            if [ -n "$result" ]; then
                break
            fi
        done
    fi
    
    # 检查测试结果
    if [ $? -ne 0 ] || [ -z "$result" ]; then
        log "ERROR" "代理测试失败（所有测试地址均无法访问）"
        echo -e "${RED}错误：代理连接测试失败${NC}"
        return 1
    else
        log "INFO" "代理测试成功，当前IP: $result"
        echo -e "${GREEN}代理测试成功！当前IP: $result${NC}"
        return 0
    fi
}

# 启动全局代理模式
start_global_proxy() {
    local proxy_port="$1"
    local redsocks_port="$2"
    
    log "INFO" "启动全局代理模式"
    
    # 配置系统代理环境变量
    configure_system_proxy "$proxy_port"
    
    # 配置并启动redsocks
    configure_redsocks "$proxy_port" "$redsocks_port"
    start_redsocks "$redsocks_port" || return 1
    
    # 配置iptables规则
    configure_iptables "$redsocks_port"
    
    log "INFO" "全局代理启动完成"
    echo -e "${GREEN}全局代理已启动！所有TCP流量将通过代理${NC}"
}

# 停止全局代理模式
stop_global_proxy() {
    log "INFO" "停止全局代理模式"
    
    # 清除iptables规则
    clear_iptables
    
    # 停止redsocks服务
    stop_redsocks
    
    # 清除系统代理环境变量
    clear_system_proxy
    
    log "INFO" "全局代理已停止"
    echo -e "${GREEN}全局代理已停止${NC}"
}

# 检查代理状态（修复：根据代理模式选择对应的IP检测方式）
check_status() {
    local proxy_port="$1"
    local redsocks_port="$2"
    local proxy_mode="$3"
    
    echo -e "${BLUE}==== 代理状态检查 ====${NC}"
    
    # 配置文件状态
    if [ -f "$CONFIG_FILE" ]; then
        echo -e "${GREEN}配置文件: 已存在${NC}"
    else
        echo -e "${RED}配置文件: 不存在${NC}"
    fi
    
    # SSH隧道状态（通过端口监听判断）
    if ss -ltn | grep -q ":$proxy_port "; then
        echo -e "${GREEN}SSH隧道: 运行中 (端口: $proxy_port)${NC}"
    else
        echo -e "${RED}SSH隧道: 未运行${NC}"
    fi
    
    # redsocks状态
    if pgrep -f "redsocks -c $REDSOCKS_CONF" >/dev/null; then
        echo -e "${GREEN}redsocks: 运行中 (端口: $redsocks_port)${NC}"
    else
        echo -e "${YELLOW}redsocks: 未运行${NC}"
    fi
    
    # iptables规则状态
    if iptables -t nat -L REDSOCKS -n 2>/dev/null | grep -q "REDIRECT"; then
        echo -e "${GREEN}iptables规则: 已配置${NC}"
    else
        echo -e "${YELLOW}iptables规则: 未配置${NC}"
    fi
    
    # 系统代理环境变量状态
    if [ -n "$http_proxy" ] || [ -f "$HOME/.proxyrc" ]; then
        echo -e "${GREEN}系统代理变量: 已设置${NC}"
    else
        echo -e "${YELLOW}系统代理变量: 未设置${NC}"
    fi
    
    # 当前代理模式
    local mode_text="未知模式"
    if [ "$proxy_mode" = "global" ]; then
        mode_text="全局代理模式"
    elif [ "$proxy_mode" = "proxychains" ]; then
        mode_text="proxychains模式"
    fi
    echo -e "${BLUE}当前代理模式: $mode_text${NC}"
    
    # 当前公网IP（修复：根据代理模式选择检测方式）
    echo -n -e "${BLUE}当前公网IP: ${NC}"
    if [ "$proxy_mode" = "proxychains" ]; then
        # proxychains模式下，通过proxychains获取代理IP
        proxy_ip=$(proxychains4 curl -s --connect-timeout 3 https://ipinfo.io/ip 2>/dev/null)
        direct_ip=$(curl -s --connect-timeout 3 https://ipinfo.io/ip 2>/dev/null)
        
        if [ -n "$proxy_ip" ] && [ "$proxy_ip" != "$direct_ip" ]; then
            echo -e "${GREEN}$proxy_ip${NC} (代理IP)"
        else
            echo -e "${YELLOW}$direct_ip${NC} (直连IP，代理可能未生效)"
        fi
    else
        # 全局模式下，直接获取IP
        echo -e "${BLUE}$(curl -s --connect-timeout 3 https://ipinfo.io/ip || echo "获取失败")${NC}"
    fi
    
    echo -e "${BLUE}=====================${NC}"
}

# 显示帮助信息
show_help() {
    echo -e "${GREEN}SSH隧道代理脚本${NC}"
    echo "用法: $0 [选项]"
    echo "选项:"
    echo "  -s, --start [模式]   启动代理（模式：global/全局代理 或 proxychains/局部代理，默认proxychains）"
    echo "  -S, --stop          停止代理"
    echo "  -r, --restart [模式] 重启代理（模式同上）"
    echo "  -c, --config        配置SSH服务器信息"
    echo "  -t, --test          测试代理连接"
    echo "  -m, --mode          切换默认代理模式"
    echo "  -C, --status        查看当前代理状态"
    echo "  -h, --help          显示帮助信息"
}

# 清理函数（用于信号捕获）
cleanup() {
    log "INFO" "接收到终止信号，开始清理代理资源..."
    # 仅在已加载配置时执行清理
    if [ -f "$CONFIG_FILE" ]; then
        load_config >/dev/null 2>&1  # 静默加载配置
        stop_global_proxy
        stop_tunnel "$PROXY_PORT"
    fi
    log "INFO" "脚本已终止，资源清理完成"
    exit 0
}

# 主函数
main() {
    # 注册信号捕获（优化：脚本终止时自动清理）
    trap cleanup SIGINT SIGTERM
    
    # 检查root权限
    check_root
    
    # 初始化日志文件
    > "$LOG_FILE"
    
    # 安装依赖
    install_dependencies
    
    # 解析命令行参数
    local ACTION=""
    local PROXY_MODE=""
    
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -s|--start)
                ACTION="start"
                if [ -n "$2" ]; then
                    case "$2" in
                        global|全局代理)
                            PROXY_MODE="global"
                            shift
                            ;;
                        proxychains|局部代理)
                            PROXY_MODE="proxychains"
                            shift
                            ;;
                        *)
                            echo -e "${RED}错误：未知模式 '$2'${NC}"
                            show_help
                            exit 1
                            ;;
                    esac
                fi
                ;;
            -S|--stop)
                ACTION="stop"
                ;;
            -r|--restart)
                ACTION="restart"
                if [ -n "$2" ]; then
                    case "$2" in
                        global|全局代理)
                            PROXY_MODE="global"
                            shift
                            ;;
                        proxychains|局部代理)
                            PROXY_MODE="proxychains"
                            shift
                            ;;
                        *)
                            echo -e "${RED}错误：未知模式 '$2'${NC}"
                            show_help
                            exit 1
                            ;;
                    esac
                fi
                ;;
            -c|--config)
                ACTION="config"
                ;;
            -t|--test)
                ACTION="test"
                ;;
            -m|--mode)
                ACTION="set_mode"
                ;;
            -C|--status)
                ACTION="status"
                ;;
            -h|--help)
                show_help
                exit 0
                ;;
            *)
                echo -e "${RED}错误：未知选项 '$1'${NC}"
                show_help
                exit 1
                ;;
        esac
        shift
    done
    
    # 配置操作
    if [ "$ACTION" = "config" ]; then
        echo -e "${GREEN}==== 配置SSH服务器信息 ====${NC}"
        read -p "SSH服务器地址: " HOST
        read -p "SSH服务器端口 [22]: " PORT
        PORT=${PORT:-22}
        read -p "SSH用户名: " USER
        
        # 解析服务器IP（优化：区分IP和域名解析）
        if [[ "$HOST" =~ ^[0-9]+\.[0-9]+\.[0-9]+\.[0-9]+$ ]]; then
            # HOST是IP地址，直接使用
            SSH_SERVER_IP="$HOST"
        else
            # HOST是域名，尝试解析IP
            SSH_SERVER_IP=$(getent hosts "$HOST" | awk '{print $1}')
        fi
        
        if [ -z "$SSH_SERVER_IP" ]; then
            echo -e "${YELLOW}警告：无法解析服务器IP，将跳过SSH服务器流量排除规则${NC}"
            SSH_SERVER_IP=""
        else
            echo -e "${YELLOW}服务器IP地址: $SSH_SERVER_IP${NC}"
        fi
        
        # 自动选择可用端口
        PROXY_PORT=$(find_available_port 10000 10100)
        echo -e "${YELLOW}已选择代理端口: $PROXY_PORT${NC}"
        
        REDSOCKS_PORT=$(find_available_port 10101 10200)
        echo -e "${YELLOW}已选择redsocks端口: $REDSOCKS_PORT${NC}"
        
        # 选择默认代理模式
        echo -e "${BLUE}请选择默认代理模式：${NC}"
        select mode in "全局代理" "proxychains模式"; do
            case $mode in
                "全局代理")
                    PROXY_MODE="global"
                    break
                    ;;
                "proxychains模式")
                    PROXY_MODE="proxychains"
                    break
                    ;;
                *)
                    echo -e "${RED}无效选择，请重新输入${NC}"
                    ;;
            esac
        done
        
        # 保存配置
        save_config "$HOST" "$PORT" "$USER" "$PROXY_PORT" "$REDSOCKS_PORT" "$SSH_SERVER_IP" "$PROXY_MODE"
        
        # 配置proxychains
        configure_proxychains "$PROXY_PORT"
        
        echo -e "${GREEN}配置完成！可使用 --start 启动代理${NC}"
        exit 0
    fi
    
    # 切换代理模式
    if [ "$ACTION" = "set_mode" ]; then
        if ! load_config; then
            echo -e "${RED}未找到配置文件，请先使用 --config 配置${NC}"
            exit 1
        fi
        
        echo -e "${BLUE}请选择默认代理模式：${NC}"
        select mode in "全局代理" "proxychains模式"; do
            case $mode in
                "全局代理")
                    NEW_MODE="global"
                    break
                    ;;
                "proxychains模式")
                    NEW_MODE="proxychains"
                    break
                    ;;
                *)
                    echo -e "${RED}无效选择，请重新输入${NC}"
                    ;;
            esac
        done
        
        # 保存新模式
        save_config "$HOST" "$PORT" "$USER" "$PROXY_PORT" "$REDSOCKS_PORT" "$SSH_SERVER_IP" "$NEW_MODE"
        
        echo -e "${GREEN}默认代理模式已设置为: $NEW_MODE${NC}"
        exit 0
    fi
    
    # 检查状态
    if [ "$ACTION" = "status" ]; then
        if ! load_config; then
            echo -e "${RED}未找到配置文件，请先使用 --config 配置${NC}"
            exit 1
        fi
        
        check_status "$PROXY_PORT" "$REDSOCKS_PORT" "$PROXY_MODE"
        exit 0
    fi
    
    # 加载配置（其他操作需要配置文件）
    if ! load_config; then
        echo -e "${RED}未找到配置文件，请先运行 --config 进行配置${NC}"
        exit 1
    fi
    
    # 如果未指定模式，使用配置中的默认模式
    if [ -z "$PROXY_MODE" ]; then
        PROXY_MODE="$PROXY_MODE"
    fi
    
    # 执行用户操作
    case "$ACTION" in
        start)
            # 启动SSH隧道
            start_tunnel "$HOST" "$PORT" "$USER" "$PROXY_PORT" || exit 1
            
            # 根据模式启动对应代理
            if [ "$PROXY_MODE" = "global" ]; then
                start_global_proxy "$PROXY_PORT" "$REDSOCKS_PORT" || exit 1
            else
                echo -e "${GREEN}proxychains代理已启动！${NC}"
                echo -e "${YELLOW}使用方法: proxychains4 [命令]${NC}"
            fi
            
            # 测试代理连接
            test_proxy "$PROXY_PORT" "$PROXY_MODE"
            ;;
        
        stop)
            # 停止全局代理（如果运行中）
            stop_global_proxy
            
            # 停止SSH隧道
            stop_tunnel "$PROXY_PORT"
            
            echo -e "${GREEN}所有代理服务已停止${NC}"
            ;;
        
        restart)
            # 先停止现有服务
            stop_global_proxy
            stop_tunnel "$PROXY_PORT"
            
            # 启动SSH隧道
            start_tunnel "$HOST" "$PORT" "$USER" "$PROXY_PORT" || exit 1
            
            # 启动对应代理模式
            if [ "$PROXY_MODE" = "global" ]; then
                start_global_proxy "$PROXY_PORT" "$REDSOCKS_PORT" || exit 1
            else
                echo -e "${GREEN}proxychains代理已重启！${NC}"
            fi
            
            # 测试代理
            test_proxy "$PROXY_PORT" "$PROXY_MODE"
            ;;
        
        test)
            test_proxy "$PROXY_PORT" "$PROXY_MODE"
            ;;
        
        *)
            show_help
            exit 1
            ;;
    esac
}

# 启动主函数
main "$@"
